
#include "libmri.h"

#include <omp.h>

#include <algorithm>
using std::max;
using std::min;

__global__ void
calib3d_cuamul (cplxf *y,
                cplxf *calib, int cs1, int cs2, int cs3, int nch,
                cplxf *x, int xs1, int xs2, int xs3)
{
	int m1 = (cs1-xs1+1), m2 = (cs2-xs1+1), m3 = (cs3-xs3+1), m = m1*m2*m3;

	for (int id = threadIdx.x + blockDim.x*blockIdx.x;
	     id < m; id  += blockDim.x*gridDim.x)
	{
		int tmp = id;
		int i3 = tmp / (m1*m2); tmp -= i3*m1*m2;
		int i2 = tmp / m1;      tmp -= i2*m1;
		int i1 = tmp;

		cplxf y_i(0.0f,0.0f);

		for (int ch = 0; ch < nch; ++ch)
		{
			for (int j3 = 0; j3 < min (xs3, cs3-i3); ++j3)
			  for (int j2 = 0; j2 < min (xs2, cs2-i2); ++j2)
				for (int j1 = 0; j1 < min (xs1, cs1-i1); ++j1)
				{
					y_i += calib [i1+j1 + cs1*(i2+j2 + cs2*(i3+j3 + cs3*ch))]
					         * x [  j1  + xs1*(  j2  + xs2*(  j3  + xs3*ch))]; 
				}
		}

		y [i1 + m1*(i2 + m2*i3)] = y_i;
	}
}

__global__ void
calib3d_cuatmul (cplxf *z, cplxf *y, cplxf *calib,
                 int cs1, int cs2, int cs3, int nch, int xs1, int xs2, int xs3)
{
	int m1 = (cs1-xs1+1), m2 = (cs2-xs1+1), m3 = (cs3-xs3+1);

	for (int id = threadIdx.x + blockDim.x*blockIdx.x;
	     id < xs1*xs2*xs3*nch; id += blockDim.x*gridDim.x)
	{
		int tmp = id;
		int ch = tmp / (xs1*xs2*xs3); tmp -= ch*xs1*xs2*xs3;
		int i3 = tmp / (xs1*xs2);     tmp -= i3*xs1*xs2;
		int i2 = tmp / xs1;           tmp -= i2*xs1;
		int i1 = tmp;

		cplxf z_i(0.0f,0.0f);

		for (int j3 = 0; j3 < min (m3, cs3-i3); ++j3)
		  for (int j2 = 0; j2 < min (m2, cs2-i2); ++j2)
			for (int j1 = 0; j1 < min (m1, cs1-i1); ++j1)
			{
				z_i += conj (calib [i1+j1 + cs1*(i2+j2 + cs2*(i3+j3 + cs3*ch))])
				               * y [  j1  +  m1*(  j2  +  m2*( j3 ))];
			}

		z [i1 + xs1*(i2 + xs2*(i3 + xs3*ch))] = z_i;
	}
}

void
calib3d_cuatamul (cplxf *z, cplxf *calib, int *csize, cplxf *x, int *xsize)
{
	int cs1 = csize[0], cs2 = csize[1], cs3 = csize[2], nch = csize[3];
	int xs1 = xsize[0], xs2 = xsize[1], xs3 = xsize[2];
	int m1 = (cs1-xs1+1), m2 = (cs2-xs1+1), m3 = (cs3-xs3+1), m = m1*m2*m3;

	int dev;
	cudaDeviceProp prop;
	cuda (GetDevice (&dev));
	cuda (GetDeviceProperties (&prop, dev));
	
	cplxf *y; cuda (Malloc (&y, m*sizeof(*y)));

	calib3d_cuamul <<<2*prop.multiProcessorCount, 64>>>
		(y, calib, cs1, cs2, cs3, nch, x, xs1, xs2, xs3);

	calib3d_cuatmul <<<2*prop.multiProcessorCount, 64>>>
	 	(z, y, calib, cs1, cs2, cs3, nch, xs1, xs2, xs3);

	cuda (Free (y));
}

#ifdef _MAIN_

int main (int argc, char **argv)
{
	shm_array_t *shm_calib = 0, *shm_x = 0;

	struct arg *args[] = {
		new arg_t <shm_array_t*> ("calib", shm_calib),
		new arg_t <shm_array_t*> ("x", shm_x),
		0};

	get_args (argc, argv, args);

	if (shm_calib == 0){
		if (get_shm_array ("calib", &shm_calib)){
			perror("calib");
			return -1;
		}
	}
	cplxf *calib = (cplxf*) shm_calib->data.c;
	int cs1 = shm_calib->dims[0],
	    cs2 = shm_calib->dims[1],
	    cs3 = shm_calib->dims[2],
	    nch = shm_calib->dims[3];

	int csize[4] = {cs1, cs2, cs3, nch};

	if (shm_x == 0){
		if (get_shm_array ("x", &shm_x)){
			perror ("x");
			return -1;
		}
	}

	int xs1 = shm_x->dims[0],
	    xs2 = shm_x->dims[1],
	    xs3 = shm_x->dims[1];

	int xsize[3] = {xs1, xs2, xs3};
	cplxf *x = (cplxf*) shm_x->data.c;

	shm_array_t z_size, *shm_z = 0;
	z_size.type = shm_array_cplxf;
	z_size.ndims = 4;
	z_size.dims[0] = xs1;
	z_size.dims[1] = xs2;
	z_size.dims[2] = xs3;
	z_size.dims[3] = nch;

	if (make_shm_array ("z", &z_size, &shm_z)){
		perror("z");
		return -1;
	}
	cplxf *z = (cplxf*) shm_z->data.c;

	calib3d_atamul (z, calib, csize, x, xsize);

	return 0;
}
#endif
